from pydantic import BaseModel
from sqlalchemy import create_engine, text
from sqlalchemy.engine import Connection
from sqlalchemy.orm import sessionmaker, declarative_base

Base = declarative_base()


class DataBase(BaseModel):
    host: str
    port: int = 3306
    user: str
    pwd: str
    db: str


def get_db_connection(database: DataBase) -> Connection:
    # database_url = "sqlite:///example.db"
    # database_url = "mysql+pymysql://root:123456@192.168.1.5:3306/db_test?charset=utf8"
    try:
        database_url = f"mysql+pymysql://{database.user}:{database.pwd}@{database.host}:{database.port}/{database.db}?charset=utf8"
        engine = create_engine(database_url)
        with engine.connect() as conn:
            yield conn  # 使用生成器返回数据库连接
    except Exception as e:
        raise e


def get_db_connection01() -> Connection:
    # database_url = "sqlite:///example.db"
    # database_url = "mysql+pymysql://root:123456@192.168.1.5:3306/db_test?charset=utf8"

    # database_url = f"mysql+pymysql://{database.user}:{database.pwd}@{database.host}:{database.port}/{database.db}?charset=utf8"
    test_database_url = "mysql+pymysql://root:admin123@localhost/tortoiseorm"
    engine = create_engine(test_database_url)
    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    db = SessionLocal()
    try:
        yield db  # 使用生成器返回数据库连接

    except Exception as e:
        raise e
    finally:
        db.close()


# 查询封装--用于造数平台，多数据库操作场景
def execute_sql(session: sessionmaker, sql: str, all: bool = True):
    result = ""
    if all:
        result = session.execute(text(sql))
        rows = result.fetchall()
        columns = result.keys()
        list_of_dict = [dict(zip(columns, row)) for row in rows]
        print(f"list_of_dict:{type(list_of_dict)},list_of_dict:{list_of_dict}")
        return list_of_dict
    else:
        result = session.execute(text(sql)).first()
        return result


def add_sql(session: sessionmaker, sql: str):
    result = session.execute(text(sql))
    session.commit()
    print(f"result:{result}")
    return result


def update_sql(session: sessionmaker, sql: str):
    result = session.execute(text(sql))
    session.commit()
    print(f"result:{result}")
    return result


class TypeCast:

    def to_dict(self):  # 方法一，该方法直接获取数据库原始数值,对于一些特殊字符如时间戳无法转换
        return {c.name: getattr(self, c.name, None) for c in self.__table__.columns}  # 记得加None(网上一些教程没有加None是无法使用的)

    Base.to_dict = to_dict  # 如果使用的是flask-sqlalchemy，就使用对应的基类

    def to_dict(self):  # 方法二，该方法可以将获取结果进行定制，例如如下是将所有非空值输出成str类型
        result = {}
        for key in self.__mapper__.c.keys():
            if getattr(self, key) is not None:
                result[key] = str(getattr(self, key))
            else:
                result[key] = getattr(self, key)
        return result

    # def to_dict(self):  # 方法二定制，将时间戳值转为str类型，其他直接输出
    # result = {}
    # for key in self.__mapper__.c.keys():
    #     if type(getattr(self, key)) == datetime.datetime:
    #         result[key] = str(getattr(self, key))
    #     else:
    #         result[key] = getattr(self, key)
    # return result

    # 配合to_dict一起使用
    def to_json(self, all_vendors):  # 多条结果时转为list(json)
        v = [ven.to_dict() for ven in all_vendors]
        return v
